package com.om.abom.utils;

import com.jcraft.jsch.*;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Properties;
import java.util.function.Function;

import static com.om.abom.config.AbomConfig.*;

@Slf4j
@Getter
public class JschExecutor
{
    private final Session session;
    private Channel channel;
    private final String ip;

    public JschExecutor(String ip) throws JSchException
    {
        this.ip = ip;
        JSch jsch = new JSch();
        session = jsch.getSession(OS_USERNAME, ip, OS_SSH_PORT);
        session.setPassword(OS_PASSWORD);
        session.setTimeout(1000);

        // 避免第一次连接时检查主机密钥的提示
        Properties config = new Properties();
        config.put("StrictHostKeyChecking", "no");
        session.setConfig(config);

        session.connect();
    }

    public String execCmd(String command) throws JSchException, IOException
    {
        // 每次任务开始前重新打开一个channel，使用旧的 channel 会阻塞
        channel = session.openChannel("exec");
        ((ChannelExec) channel).setCommand(command);
        channel.setInputStream(null);
        ((ChannelExec) channel).setErrStream(System.err);
        // 执行命令
        channel.connect();
        // 读取执行结果
        String output = processStream(channel.getInputStream());
        channel.disconnect();
        return output;
    }

    public static String processStream(InputStream in)
    {
        StringBuilder stringBuilder = new StringBuilder();
        try {
            BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(in, DEFAULT_CHARSET));
            String line;
            while ((line = bufferedReader.readLine()) != null) {
                stringBuilder.append(line).append("\n");
            }
        } catch (IOException e) {
            log.error("processStream error", e);
        }
        return stringBuilder.toString();
    }

    // 多次执行命令
    public static String execMultiCmd(String ip, Function<JschExecutor, String> doExecMultiCmd)
    {
        JschExecutor executor = null;
        String errorMsg = null;
        try {
            executor = new JschExecutor(ip);
            errorMsg = doExecMultiCmd.apply(executor);
        } catch (JSchException e) {
            errorMsg = ip + " 主机连接不上";
            log.error(errorMsg, e);
        } finally {
            // 在一个 channel 执行多条命令，在这里统一关闭连接
            if (executor != null) {
                executor.disconnect();
            }
        }
        return errorMsg;
    }

    public void disconnect()
    {
        log.warn("disconnect {}", ip);
        channel.disconnect();
        session.disconnect();
    }
}
